import backtrader as bt
from backtrader import cerebro

import pandas as pd
import numpy as np
from datetime import timedelta, datetime
from tabulate import tabulate

import tushare as ts
ts.set_token('ec3a7bd2ac0ba626ce4bf97a36feb38653b15fc379f3e3bcd267644a')
pro = ts.pro_api()

#正常显示画图时出现的中文和负号
from pylab import mpl
mpl.rcParams['font.sans-serif']=['SimHei']
mpl.rcParams['axes.unicode_minus']=False

g_init_cash = 100000.0

def prepare_data(start_date, end_date):
    start_date_short, end_date_short = start_date.replace('-', ''), end_date.replace('-', '')

    # 获取回测区间内还在交易的可转债基础数据列表
    df_bond = pro.cb_basic(tscode='128016.SZ', fields="ts_code,bond_short_name,stk_code,stk_short_name,list_date,delist_date,conv_price")
    # 过滤掉: 未上市(list_date==None), 必须回测日期内上市的(list_date<end_date), 下市日期>start_date, 未下市
    df_bond = df_bond[( (~df_bond['list_date'].isnull()) & (df_bond['list_date'] < end_date)) 
                            & ( (df_bond['delist_date'] > start_date) | (df_bond['delist_date'].isnull()) )]
    df_bond = df_bond.rename(columns={"conv_price":"bond_conv_price"}).reset_index(drop=True)
    # 测试，只取前N个                        
    # df_bond = df_bond.head(80)
    print('回测区间内共计找到%d只可转债' % df_bond.shape[0])

    # 获取对应的可转债的日线行情数据
    bond_data = dict()
    for i, bond_row in df_bond.iterrows():
        print('%d 准备数据: [%s(%s)]' % (i+1, bond_row['bond_short_name'], bond_row['ts_code']))

        # 取可转债行情数据 (数据需要按日期正向排序)
        df_bond_bar = pro.cb_daily(ts_code=bond_row['ts_code'], start_date=start_date_short, end_date=end_date_short).sort_values(by='trade_date', ascending=True)
        if df_bond_bar.empty:
            print('[%s(%s)]没有行情数据，忽略' % (bond_row['bond_short_name'], bond_row['ts_code']))
            continue

        # 将open=0的行情数据清洗
        df_bond_bar = df_bond_bar[df_bond_bar['open'] != 0]

        if len(df_bond_bar) < 30:
            print('[%s(%s)]行情数据过少，忽略' % (bond_row['bond_short_name'], bond_row['ts_code']))
            continue

        df_bond_bar = pd.merge(df_bond, df_bond_bar, on=['ts_code'])
        df_bond_bar = df_bond_bar.rename(columns={"ts_code":"bond_code", "open":"bond_open", "high":"bond_high", "low":"bond_low", "close":"bond_close", "vol":"bond_vol"})
        # print(df_bond_bar.head(10))

        # 取对应正股的当日的股票行情数据
        df_staock_bar = pro.daily(ts_code=bond_row['stk_code'], start_date=start_date_short, end_date=end_date_short).sort_values(by='trade_date', ascending=True)
        df_staock_bar = df_staock_bar.rename(columns={"ts_code":"stk_code", "open":"stock_open", "high":"stock_high", "low":"stock_low", "close":"stock_close", "vol":"stock_vol", "amount":"stock_amount",
                                                        "pre_close":"stock_pre_close", "change":"stock_change", "pct_chg":"stock_pct_chg"})
        # 和可转债行情数据合并
        df = pd.merge(df_bond_bar, df_staock_bar, on=['stk_code', 'trade_date'])
        df = df[['trade_date', 'bond_code', 'bond_short_name', 'stk_code', 'stk_short_name', 'bond_conv_price', 'bond_open', 'bond_high', 'bond_low', 'bond_close', 'bond_vol', 'stock_open', 'stock_high', 'stock_low', 'stock_close']]

        # 计算可转债每日的双低数据
        df['转股价值'] = 100 / df['bond_conv_price'] * df['stock_close']
        df['溢价率'] = (df['bond_close'] /df['转股价值'] -1) * 100
        df['双低'] = df['转股价值'] + df['溢价率']
        
        df.index = pd.to_datetime(df.trade_date)
        # print(df.head(10))

        # bond_code => dataframe
        bond_data[bond_row['bond_short_name'] + '(' + bond_row['ts_code'] + ')'] = df

    # print(bond_data.get('113009.SH'))
    return bond_data

# Trade list similar to Amibroker output
class TradeAnalyzer(bt.Analyzer):
    """
    ref - bt's unique trade identifier
    ticker - data feed name
    datein - date and time of trade opening
    pricein - price of trade entry                      买入价
    dir - long or short
    dateout - date and time of trade closing
    priceout - price of trade exit                      卖出价
    chng% - exit price to entry price ratio             卖出价相对买入价的盈利百分比
    pnl - money profit/loss per trade                   盈亏
    pnl% - proft/loss in %s to broker's value at the trade closing  盈亏占交易关闭时账户总值的百分比
    cumpnl - cumulative profit/loss                     累计的盈亏
    size - max position size during trade               交易中最大的成交数量
    value - max trade value                             交易中最大的成交总值
    nbars - trade duration in bars                      交易持续的bar数
    pnl/bar - profit/loss per bar                       平均每根bar的盈亏
    mfe% - max favorable excursion                      最大有利变动幅度
    mae% - max adverse excursion                        最大不利变动幅度
    """
    def get_analysis(self):
        return self.trades

    def __init__(self):
        self.trades = []
        self.cumprofit = 0.0

    def notify_trade(self, trade):
        if trade.isclosed:
            brokervalue = self.strategy.broker.getvalue()
            dir = 'short'
            if trade.history[0].event.size > 0: dir = 'long'

            pricein = trade.history[len(trade.history)-1].status.price
            priceout = trade.history[len(trade.history)-1].event.price
            datein = bt.num2date(trade.history[0].status.dt)
            dateout = bt.num2date(trade.history[len(trade.history)-1].status.dt)
            if trade.data._timeframe >= bt.TimeFrame.Days:
                datein = datein.date()
                dateout = dateout.date()

            pcntchange = 100 * priceout / pricein - 100
            pnl = trade.history[len(trade.history)-1].status.pnlcomm
            pnlpcnt = 100 * pnl / brokervalue
            barlen = trade.history[len(trade.history)-1].status.barlen
            pbar = pnl / barlen
            self.cumprofit += pnl

            size = value = 0.0
            for record in trade.history:
                if abs(size) < abs(record.status.size):
                    size = record.status.size
                    value = record.status.value

            highest_in_trade = max(trade.data.high.get(ago=0, size=barlen+1))
            lowest_in_trade = min(trade.data.low.get(ago=0, size=barlen+1))
            hp = 100 * (highest_in_trade - pricein) / pricein
            lp = 100 * (lowest_in_trade - pricein) / pricein
            if dir == 'long':
                mfe = hp
                mae = lp
            if dir == 'short':
                mfe = -lp
                mae = -hp

            self.trades.append({'ref': trade.ref, 'ticker': trade.data._name, 'dir': dir,
                 'datein': datein, 'pricein': pricein, 'dateout': dateout, 'priceout': priceout,
                 'chng%': round(pcntchange, 2), 'pnl': pnl, 'pnl%': round(pnlpcnt, 2),
                 'size': size, 'value': value, 'cumpnl': self.cumprofit,
                 'nbars': barlen, 'pnl/bar': round(pbar, 2),
                 'mfe%': round(mfe, 2), 'mae%': round(mae, 2)})

class BondData(bt.feeds.PandasData):
    lines = ('doublelow', ) # 双低 line
    params = dict(
        doublelow=-1,
        dtformat = '%Y-%m-%d',
        openinterest = -1
    )

class BondStrategy(bt.Strategy):
    params = dict(
        max_hold_num = 20,       # 最多持有可转债数量
        stop_loss = 5,         # 亏损n%止损
        take_profit = 10,       # 获利n%止盈

        limit_percent = 1,     # limit 1% percent for order
        trail_percent = 1,     # stop-trail percent 3% for order

        ma_fast = 5,
        ma_slow1 = 10,
        ma_slow2 = 20,
    )

    def __init__(self):
        self.bonds = self.datas[1:] # 去掉第一个指数
        self.hold_bonds = [] # 目前持有的可转债
        self.orders = []    # 所有订单列表
        self.ranks = []     # 可转债评分的列表
        self.fee = 0.0      # 合计费用 (佣金+印花税)
        self.exclude_bonds = [] # 排除的标的 (之前已经亏损达到止损，后续排除)
        self.ind_buy, self.ind_sell = dict(), dict()       # indicator

        for data in self.datas:
            ma_fast = bt.ind.SMA(data.close, period=self.p.ma_fast)
            ma_slow1 = bt.ind.SMA(data.close, period=self.p.ma_slow1)
            ma_slow2 = bt.ind.SMA(data.close, period=self.p.ma_slow2)
            # self.ind_buy[data] = bt.ind.CrossOver(ma_fast, ma_slow1)
            # 买入信号: 多头排列
            self.ind_buy[data] = bt.And(ma_fast > ma_slow1, ma_slow1 > ma_slow2)
            # 卖出信号: 5日均线交叉20日均线向下
            self.ind_sell[data] = bt.ind.CrossDown(ma_fast, ma_slow2)

    def log(self, arg):
        # print('{} {}'.format(self.datetime.date(), arg))
        pass

    def prenext(self):
        self.log('prenext()')
        self.next()

    def nextstart(self):
        self.log('nextstart()')
        self.next()

    def next(self):
        self.curDate = self.data0.datetime.date(0)
        self.log('-------------------------- next -------------------------')

        # for o in self.orders:
        #     self.log('当前订单: [%s][%s]' % (o.data._name, o.status))

        self.hold_bonds = [b for b, pos in self.getpositions().items() if pos]

        # STEP1: 先检查是否有需要卖出的可转债，将资金空出来
        self.log('总资产:%.2f元(空闲:%.2f元), 总费用:%.2f元' % (self.broker.getvalue(), self.broker.getcash(), self.fee))
        for bond in self.hold_bonds:
            if self.getposition(bond).size > 0:
                # 持仓成本
                cost = self.getposition(bond).size * self.getposition(bond).price
                # 当前市值
                current_value = self.getposition(bond).size * bond.close[0]
                # 持仓利润
                profit = current_value - cost
                profit_percent = (profit/cost)*100.0

                self.log('持有: [%s] %.2f元(%d股x%.2f元) 盈亏:%.2f元(%.2f%%)' % 
                    (bond._name, 
                    cost,
                    self.getposition(bond).size, 
                    self.getposition(bond).price,
                    profit,
                    profit_percent))

                # 开始检查是否需要止盈或止损
                found_order = bond._name in [o.data._name for o in self.orders]
                if not found_order: #确保订单没有正在处理
                    # 止盈用跟踪止盈
                    if profit > 0 and profit_percent > self.p.take_profit:
                        self.log('- 调出(止盈): [%s], 盈亏:%.2f元(%.2f%%)' % (bond._name, profit, profit_percent))
                        order = self.sell(exectype=bt.Order.StopTrail,
                                                data=bond,
                                                size=self.getposition(bond).size,
                                                trailpercent=self.p.trail_percent/100.0)
                        self.orders.append(order)

                    # 止损的时候麻溜点
                    if profit < 0 and abs(profit_percent) > self.p.stop_loss:
                        self.log('- 调出(止损): [%s], 盈亏:%.2f元(%.2f%%)' % (bond._name, profit, profit_percent))
                        order = self.sell(data=bond,
                                            size=self.getposition(bond).size)
                        self.orders.append(order)
                        self.exclude_bonds.append(order.data._name)

                    if self.ind_sell[bond][0] < 0:
                        self.log('- 调出(指标反转): [%s], 盈亏:%.2f元(%.2f%%)' % (bond._name, profit, profit_percent))
                        order = self.sell(data=bond,
                                            size=self.getposition(bond).size)
                        self.orders.append(order)
                        self.exclude_bonds.append(order.data._name)


        # STEP2: 按照买入的算法排序，买入可转债标的

        # 按双低价格排序
        self.bonds.sort(key=lambda bond:bond.doublelow[0])
        # 取指定的最低的几只可转债
        # self.ranks = self.bonds[0:self.p.max_hold_num]
        self.ranks = self.bonds #不截取了，后面会检查是否超过最大持仓的数量
        # for b in self.ranks:
        #     self.log('[%s]双低值:%.2f' % (b._name, b.doublelow[0]))

        # 本次选中的标的下单(必须10x倍), 资金分成max_hold_num份，每份资金按照当日收盘价计算可购买的可转债数量(10的整数倍)
        for b in (set(self.ranks) - set(self.hold_bonds)):
            # 要买入的标的简单的判断一下，需要快慢均线交叉向上
            if self.ind_buy[b][0] == None or self.ind_sell[b][0] == None:
                self.log('# [%s]暂未生成指标值, next...' % (b._name,))
                return

            #todo 不知道为什么双低能选出来这么高价格的可转债, 先限制一下120
            if self.ind_buy[b][0] > 0 and b.open[0] < 120: 
                if len(self.hold_bonds) >= self.p.max_hold_num:
                    self.log('当前持仓达到上限，不进行调仓操作，耐心等待平仓')
                    return

                # 没有正在处理的订单，不在黑名单中就下单吧
                found_order = b._name in [o.data._name for o in self.orders]
                if (not found_order) and (not b._name in self.exclude_bonds):
                    target_money = (self.broker.getvalue()*0.95) / self.p.max_hold_num
                    target_size = int(target_money / b.close[0] / 10)
                    limit_price = b.close[0] * (1.0 - self.p.limit_percent / 100.0)
 
                    self.log('+ 调入: [%s], %.2f元x%d股=%.2f元, 买入限价:%.2f元' % (b._name, b.close[0], target_size*10, b.close[0]*target_size*10, limit_price))
                    # limit买入单
                    # order = self.buy(exectype=bt.Order.Limit,
                    #                     data=b,
                    #                     size=target_size*10,
                    #                     price=limit_price,
                    #                     valid=b.datetime.date(0) + timedelta(days=3))
                    # stoplimit单                    
                    # order = self.buy(exectype=bt.Order.StopLimit,
                    #                     data=b,
                    #                     size=target_size*10,
                    #                     price=self.data.close[0],
                    #                     plimit=self.data.close[0] * (1 + self.p.limit_percent / 100.0),
                    #                     valid=b.datetime.date(0) + timedelta(days=3))
                    # 普通市价单
                    order = self.buy(data=b, size=target_size*10)
                    self.orders.append(order)

    def stop(self):
        self.log('-------------------------- stop -------------------------')
        self.log('最终总资产:%.2f元, 盈利:%.2f元, 总费用:%.2f元' % (self.broker.getvalue(), self.broker.getvalue() - g_init_cash, self.fee))

        print('-------------------------- stop -------------------------')
        print('take_profile=%d, stop_lost=%d trail_percent=%d' % (self.p.take_profit, self.p.stop_loss, self.p.trail_percent))
        print('最终总资产:%.2f元, 盈利:%.2f元, 总费用:%.2f元' % (self.broker.getvalue(), self.broker.getvalue() - g_init_cash, self.fee))


    def notify_order(self, order):
        self.log('-------------------------- notify_order -------------------------')
        for o in self.orders:
            if o.isbuy():
                self.log('当前订单: [买入][%s][%s]' % (o.data._name, o.status))
            else:
                self.log('当前订单: [卖出][%s][%s]' % (o.data._name, o.status))

        if order.status == order.Created:
            self.log('订单创建: [%s]' % order.data._name)
        elif order.status == order.Submitted:
            self.log('订单提交: [%s]' % order.data._name)
            if order.data.close[0] <= 0:
                self.log('! 订单需要被取消:[%s] (无行情数据?)' % order.data._name)
                self.cancel(order)
                need_remove = True
        elif order.status == order.Accepted:
            self.log('订单接受: [%s]' % order.data._name)


        need_remove = False
        if order.status in [order.Completed]:
            self.fee += order.executed.comm
            if order.isbuy():
                self.log(
                    '订单买入: [%s] 买入价:%.2f元, 总价:%.2f元, 费用:%.2f元' %
                    (order.data._name,  
                    order.executed.price,
                    order.executed.value,
                    order.executed.comm))
            else:  # Sell
                self.log('订单卖出: [%s] 卖出价:%.2f元, 总价:%.2f, 费用:%.2f元' %
                    (order.data._name, 
                    order.executed.price,
                    order.executed.value,
                    order.executed.comm))
            # self.bar_executed = len(self)
            need_remove = True

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消(拒绝): [%s]' % order.data._name)
            need_remove = True

        elif order.status in [order.Expired]:
            self.log('订单过期: [%s]' % order.data._name)
            need_remove = True

        if need_remove:
            # 订单‘部分’执行的情况还需要继续追踪
            for o in self.orders:
                if o.data._name == order.data._name:
                    self.orders.remove(o)

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log('交易完成: [%s] 毛收益:%.2f元, 净收益:%.2f元' %
                 (trade.data._name, trade.pnl, trade.pnlcomm))

if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 回测起始日
    start_date, end_date = '2020-01-01', '2020-12-31'

    # 获取上证指数日线行情数据，作为后续benchmark以及驱动数据
    df_index = pro.index_daily(ts_code='000001.SH', start_date=start_date.replace('-', ''), end_date=end_date.replace('-', '')).sort_values(by="trade_date" , ascending=True)
    df_index.index = pd.to_datetime(df_index.trade_date)
    data = bt.feeds.PandasData(
        dataname=df_index,
        open=3,
        high=4,
        low=5,
        close=2,
        volume=9,
        openinterest=-1,
        fromdate=datetime(2020, 1, 1), 
        todate=datetime(2020, 12, 31))
    cerebro.adddata(data, name='index')

    # 准备可转债数据并添加
    bond_data = prepare_data(start_date, end_date)
    for bond_code, bond_df in bond_data.items():
        data = BondData(
            dataname=bond_df,
            # code='bond_code',
            open='bond_open',
            high='bond_high',
            low='bond_low',
            close='bond_close',
            volume='bond_vol',
            doublelow='双低', #双低
            openinterest=-1, # 无未平仓量
            fromdate=datetime(2020, 1, 1), 
            todate=datetime(2020, 12, 31)
        )
        cerebro.adddata(data, name=bond_code)

    cerebro.broker.setcash(g_init_cash)
    cerebro.broker.setcommission(commission=0.00025) # 万分之2.5

    cerebro.addanalyzer(TradeAnalyzer, _name='trade_list')
    # cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name = 'SharpeRatio')
    # cerebro.addanalyzer(bt.analyzers.DrawDown, _name='DrawDown')

    cerebro.addstrategy(BondStrategy)
    strats = cerebro.run(tradehistory=True)
    # get analyzers data
    trade_list = strats[0].analyzers.trade_list.get_analysis()
    print('\n' + tabulate(trade_list, headers="keys") + '\n')

    # 参数调优
    # strats = cerebro.optstrategy(
    #     BondStrategy,
    #     take_profit = [3, 5, 7, 10],
    #     stop_loss = [1, 3, 5, 10],
    #     trail_percent = [1, 2, 3])
    # cerebro.run(maxcpus=1, tradehistory=True)

    # cerebro.plot(style='candlestick')
    # cerebro.plot()